# 读取相机位姿实现计算相机位姿的功能
# 可视化相机的位姿 均匀位姿采样
# from mathutils import Euler, Matrix, Vector
#使用open3d 可视化视角
import os
import cv2
import sys
import json

import numpy as np
import open3d as o3d


sys.path.append(os.getcwd())
def get_camera_intrinsic(folder):
    with open(f'{folder}/intrinsics.json', 'r') as f:
        camera_intrinsics = json.load(f)

    K = np.zeros((3, 3), dtype='float64')
    K[0, 0], K[0, 2] = float(camera_intrinsics['fx']), float(camera_intrinsics['ppx'])
    K[1, 1], K[1, 2] = float(camera_intrinsics['fy']), float(camera_intrinsics['ppy'])

    K[2, 2] = 1.
    return (camera_intrinsics, K)

def look_at(location):
    """
    Get object pose from a viewpoint location
    # Taken from https://github.com/thodan/bop_toolkit/blob/master/bop_toolkit_lib/view_sampler.py#L216
    IMPORTANT: output of this function is the object pose defined in OPENGL coordinate convention
    """
    f = -np.array(location)  # Forward direction.
    f /= np.linalg.norm(f)

    u = np.array([0.0, 0.0, 1.0])  # Up direction.
    s = np.cross(f, u)  # Side direction.
    if np.count_nonzero(s) == 0:
        # f and u are parallel, i.e. we are looking along or against Z axis.
        s = np.array([1.0, 0.0, 0.0])
    s /= np.linalg.norm(s)
    u = np.cross(s, f)  # Recompute up.
    R = np.array([[s[0], s[1], s[2]],
                  [u[0], u[1], u[2]],
                  [-f[0], -f[1], -f[2]]])
    t = - R.dot(np.array(location).reshape((3, 1)))
    obj_pose = np.zeros((4, 4))
    obj_pose[:3, :3] = R
    obj_pose[:3, 3] = -t.reshape(-1)
    obj_pose[3, 3] = 1
    return obj_pose

def opencv2opengl(cam_matrix_world):
    """
    Change coordinate system from OpenCV to OpenGL or from OpenGL to OpenCV
    """
    from scipy.spatial.transform import Rotation as R
    rot180x = R.from_euler('x', 180, degrees=True).as_matrix()
    rotation = cam_matrix_world[:3, :3]
    translation = cam_matrix_world[:3, 3]
    # Apply the rotation to the rotation part and translation part
    output_rotation = np.dot(rot180x, rotation)
    output_translation = np.dot(rot180x, translation)
    
    # Create an output matrix of 4x4
    output_matrix = np.eye(4)
    output_matrix[:3, :3] = output_rotation
    output_matrix[:3, 3] = output_translation
    return output_matrix

# 获取相机位姿值来来自位置点

def get_camera_pose(location_txt):
    camera_poses = []
    locations = np.loadtxt(location_txt)
    locations = locations.reshape(-1, 3)

    for i in range(len(locations)):
        camera_pose = look_at(locations[i])
        camera_poses.append(camera_pose)
    return camera_poses

def vis_camera_view(n,folder,scene):
    camera_intrinsics, K = get_camera_intrinsic(folder)
    extrinsics = get_camera_pose("3d/points_2.txt")
    scene = scene.scale(10, center=scene.get_center())
    vizualizer = o3d.visualization.Visualizer()
    vizualizer.create_window(width=WIDTH, height=HEIGHT)
    coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.4, origin=[0, 0, 0])
    vizualizer.add_geometry(scene)
    vizualizer.add_geometry(coordinate_frame)

    for i in range(0, n):
        extrinsic = extrinsics[i]
        # if extrinsic[0,0]>0:
        if True:
            cameraLines = o3d.geometry.LineSet.create_camera_visualization(view_width_px=camera_intrinsics['width'],
                                                                        view_height_px=camera_intrinsics['height'],
                                                                        intrinsic=K[:3, :3], extrinsic=extrinsic, scale=0.05)
            vizualizer.add_geometry(cameraLines)


    vizualizer.run()
    return vizualizer


'''
range(0, len(extrinsics), len(extrinsics)//n)生成一个整数序列，起始值为0，终止值为len(extrinsics)（不包含），步长为len(extrinsics)//n。


'''

folder = r"E:/pose/datasets/obj_bop/obj_000000/val"
camera_intrinsics, K = get_camera_intrinsic(folder)
scene = o3d.io.read_triangle_mesh(r"E:/pose/datasets/obj_bop/obj_000000/val/model/obj_000000.ply")

WIDTH = 640
HEIGHT = 480


if __name__ == "__main__":
    location_txt = "3d/points_2.txt"
    camera_poses = get_camera_pose(location_txt)
    vizualizer=vis_camera_view(len(camera_poses),folder,scene)
    vizualizer.capture_screen_image("3d/camera_view.png")